import numpy as np
import sys
import io
import random

sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding = "utf-8")

def getWords(filepath):
    f = open(filepath, encoding = "utf-8")
    fStrings = f.readlines()
    f.close()
    replaced = [r".", r"?", r"!", r",", r"/", "\'", "\"", r"。", r"，", r"？", r"！", r":", r"：", r"；", r";"]
    for i in range(len(fStrings)):
        fStrings[i] = fStrings[i].strip()
        for item in replaced:
            fStrings[i] = fStrings[i].replace(item, " ")
        fStrings[i] = fStrings[i].split()
    return fStrings

# cnwords = getWords("cn1.txt")
# print(len(cnwords))

# enwords = getWords("en.txt")
# print(len(enwords))


class IBM_Model():

    def __init__(self, cnPath, enPath, length = 1000, epoches = 10):
        self.enWords = self.getWords(enPath)
        self.cnWords = self.getWords(cnPath)
        self.epoches = epoches
        self.length = length
        self.q = {}
        self.t = {}
        self.countWordsUnit = {}
        self.countWordsSingle = {}
        self.countPositionUnit = {}
        self.countPositionConditional = {}
        self.gama = {}
        self.bestAlignment = []
        for i in range(self.length):
            l = len(self.enWords[i])
            m = len(self.cnWords[i])
            tmpBestAlignment = [-1 for i in range(m)]
            self.bestAlignment.append(tmpBestAlignment)
            for g in range(l):
                for h in range(m):
                    self.q[str(g) + "|" + str(h) + "," + str(l) + "," + str(m)] = random.random()
                    self.t[self.cnWords[i][h] + "|" + self.enWords[i][g]] = random.random()
                    self.countWordsUnit[self.enWords[i][g] + "," + self.cnWords[i][h]] = 0.0
                    self.countWordsSingle[self.enWords[i][g]] = 0.0
                    self.countPositionUnit[str(g) + "|" + str(h) + "," + str(l) + "," + str(m)] = 0.0
                    self.countPositionConditional[str(h) + "," + str(l) + "," + str(m)] = 0.0
        print("INIT END.")
        
    def getWords(self, filepath):
        f = open(filepath, encoding = "utf-8")
        fStrings = f.readlines()
        f.close()
        replaced = [r".", r"?", r"!", r",", r"/", "\'", "\"", r"。", r"，", r"？", r"！", r":", r"：", r"；", r";"]
        for i in range(len(fStrings)):
            fStrings[i] = fStrings[i].strip()
            for item in replaced:
                fStrings[i] = fStrings[i].replace(item, " ")
            fStrings[i] = fStrings[i].split()
        return fStrings
    
    def setCountsToZero(self):
        for i in range(self.length):
            l = len(self.enWords[i])
            m = len(self.cnWords[i])
            for g in range(l):
                for h in range(m):
                    self.countWordsUnit[self.enWords[i][g] + "," + self.cnWords[i][h]] = 0.0
                    self.countWordsSingle[self.enWords[i][g]] = 0.0
                    self.countPositionUnit[str(g) + "|" + str(h) + "," + str(l) + "," + str(m)] = 0.0
                    self.countPositionConditional[str(h) + "," + str(l) + "," + str(m)] = 0.0
            print("set zero " + str(i))

    def evaluateWithEM(self):
        for s in range(self.epoches):
            self.setCountsToZero()
            print("set zero end.")
            for k in range(self.length):
                tmp = 0.0
                for i in range(len(self.cnWords[k])):
                    for j in range(len(self.enWords[k])):
                        tmp += self.q[str(j) + "|" + str(i) + "," + str(len(self.enWords[k])) + "," + str(len(self.cnWords[k]))] * self.t[self.cnWords[k][i] + "|" + self.enWords[k][j]]
                # print("calculate tmp.")
                for i in range(len(self.cnWords[k])):
                    for j in range(len(self.enWords[k])):
                        self.gama[str(k) + "," + str(i) + "," + str(j)] = self.q[str(j) + "|" + str(i) + "," + str(len(self.enWords[k])) + "," + str(len(self.cnWords[k]))] *\
                                                                          self.t[self.cnWords[k][i] + "|" + self.enWords[k][j]] / tmp 
                # print("calculate gamma.")
                for i in range(len(self.cnWords[k])):
                    for j in range(len(self.enWords[k])):
                        self.countWordsUnit[self.enWords[k][j] + "," + self.cnWords[k][i]] += self.gama[str(k) + "," + str(i) + "," + str(j)]
                        self.countWordsSingle[self.enWords[k][j]] += self.gama[str(k) + "," + str(i) + "," + str(j)] 
                        self.countPositionUnit[str(j) + "|" + str(i) + "," + str(len(self.enWords[k])) + "," + str(len(self.cnWords[k]))] += self.gama[str(k) + "," + str(i) + "," + str(j)] 
                        self.countPositionConditional[str(i) + "," + str(len(self.enWords[k])) + "," + str(len(self.cnWords[k]))] += self.gama[str(k) + "," + str(i) + "," + str(j)] 
                print("epoch: " + str(s) + " calculate counts" + str(k))
            for key, value in self.t.items():
                tmpList = key.split("|")
                if self.countWordsSingle[tmpList[1]] == 0.0:
                    self.t[key] = 0.0
                else:
                    self.t[key] = self.countWordsUnit[tmpList[1] + "," + tmpList[0]] / self.countWordsSingle[tmpList[1]]
            for key, value in self.q.items():
                tmpList = key.split("|")
                if self.countPositionConditional[tmpList[1]] == 0.0:
                    self.q[key] = 0.0
                else:
                    self.q[key] = self.countPositionUnit[key] / self.countPositionConditional[tmpList[1]]
            print("evaluate epoch:" + str(s) + " done.")

    def getBestAlignment(self):
        for k in range(self.length):
            for j in range(len(self.cnWords[k])):
                paj = 0.0
                for i in range(len(self.enWords[k])):
                    tmp = self.q[str(i) + "|" + str(j) + "," + str(len(self.enWords[k])) + "," + str(len(self.cnWords[k]))] * self.t[self.cnWords[k][j] + "|" + self.enWords[k][i]]
                    if tmp > paj:
                        paj = tmp
                        self.bestAlignment[k][j] = i
            print("GotBestAlignment for " + str(k))
    def writeBestAlignmentToFile(self, filename = "bestAlignment.txt"):
        
        f = open(filename, 'w')
        for item in self.bestAlignment:
            for i in range(len(item)):
                f.write(str(item[i]))
                f.write(" ")
            f.write("\n")
        f.close()

                    
if __name__ == '__main__':
    model = IBM_Model("cn1.txt", "en.txt", length = 100000, epoches = 100)
    model.evaluateWithEM()
    model.getBestAlignment()
    model.writeBestAlignmentToFile()